#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
import tools

# 取前50组数据方便看图
(X, Y) = tools.oneSamples(50)

# 显示样本数据
plt.plot(X[:, 1], Y, color='k', label='sample')
plt.legend()

# 最小二乘法求解 theta
theta = tools.leastSquare(X, Y)

# 求估计值
hatY = X * theta

# 画出图像
plt.plot(X[:, 1], hatY, color='r', label='least squares')
plt.legend()

# 使用梯度下降求解
th = np.matrix([[0], [0]])

# 传入不同的 alpha, 观察函数图像
(theta, J, counter) = gradientDescent(th, X, Y, alpha=1e-3)
plt.plot(X[:, 1], X * theta, color='g', label='grad-0.001')
plt.legend()

(theta, J, counter) = gradientDescent(th, X, Y, alpha=0.01)
plt.plot(X[:, 1], X * theta, label='grad-0.01')
plt.legend()

(theta, J, counter) = gradientDescent(th, X, Y, alpha=0.1)
plt.plot(X[:, 1], X * theta, label='grad-0.1')
plt.legend()

# 保存图片，非必须
plt.savefig('plot.png')
